import os,sys
import cv2
import glob
from test_ui import Ui_Dialog
from faceadd_ui import Ui_Dialog_faceadd
import numpy as np
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QVBoxLayout, QFileDialog  
from PyQt5.QtGui import QPixmap, QImage 
from PyQt5.QtCore import Qt,QTimer  
import face_recognition
import time
from PIL import Image, ImageDraw, ImageFont

g_list_faceencode = []
g_list_facename = []

def cv2ImgAddText(img, text, left, top, textColor=(0, 255, 0), textSize=20):
    if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    # 字体的格式
    fontStyle = ImageFont.truetype("simsun.ttc", textSize, encoding="utf-8")
    # 绘制文本
    draw.text((left, top), text, textColor, font=fontStyle)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
class FaceAddDlg(Ui_Dialog_faceadd):
    def __init__(self, FaceAddDlg) -> None:
        super().setupUi(FaceAddDlg)
        Dialogfaceadd.setWindowTitle("录入人脸")
        self.pushButton_add.setText("添加人脸")
        self.pushButton_exit.setText("退出")
        self.pushButton_add.clicked.connect(self.FaceAdd_clicked)
        self.pushButton_exit.clicked.connect(self.ExitDlg_clicked)
        self.faceaddlist=[]

        # 创建定时器  
        self.timer = QTimer()  
        self.timer.timeout.connect(self.update_label)  # 连接 timeout 信号到 update_label 槽函数  
        self.timer.start(10)  # 每秒（1000毫秒）触发一次 
        self.addcount = 0 
        self.progressBar_faceadding.setValue(self.addcount)

        self.label_face_show.setText("")

    def ExitDlg_clicked(self):
        Dialogfaceadd.reject()

    def FaceAdd_clicked(self):        
        folder_path = QFileDialog.getExistingDirectory(caption="选择人脸库文件夹")  
        self.faceaddlist = glob.glob(f"{folder_path}/*")
        self.progressBar_faceadding.setRange(0, len(self.faceaddlist))
        self.addcount = 0
        self.progressBar_faceadding.setValue(self.addcount)
        
        
    def update_label(self):  
        if len(self.faceaddlist) > 0:
            p_face = self.faceaddlist[0]
            self.addcount += 1
            self.progressBar_faceadding.setValue(self.addcount)
            with open(p_face, "rb") as file:  
                image_data = file.read()
            # 将字节流转换为NumPy数组  
            image_array = np.frombuffer(image_data, np.uint8)  
            image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)  
            if image is not None: 
                img_encoding = face_recognition.face_encodings(image)[0]
                if len(img_encoding) != 128:
                    return
                #print(img_encoding)
                g_list_faceencode.append(img_encoding)
                name = os.path.splitext(os.path.basename(p_face))[0]
                g_list_facename.append(name)
                face_locations = face_recognition.face_locations(image)

                for i in range(len(face_locations)):
                    cv2.rectangle(image,(face_locations[i][1],face_locations[i][0]),(face_locations[i][3],face_locations[i][2]),color=(0,255,0))

                rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 
                h, w, ch = rgb_image.shape  
                bytes_per_line = ch * w  
                convert_to_Qt_format = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)  
                # # 将 QImage 转换为 QPixmap 并显示在 QLabel 上  
                # pixmap = QPixmap.fromImage(convert_to_Qt_format)  
                # 将 QImage 转换为 QPixmap 并保持宽高比  
                pixmap = QPixmap.fromImage(convert_to_Qt_format).scaled(  
                    self.label_face_show.size(),  
                    Qt.KeepAspectRatio,  
                    Qt.SmoothTransformation  
                ) 
                # 将 QPixmap 设置为 QLabel 的内容  
                self.label_face_show.setPixmap(pixmap) 
            self.faceaddlist = self.faceaddlist[1:]
            


class MainDlg(Ui_Dialog):
    def __init__(self, MainDlg) -> None:
        super().setupUi(MainDlg)
        self.pushButton_add.setText("注册人脸")
        self.pushButton_reg.setText("开始识别")

        self.pushButton_add.clicked.connect(self.button1_clicked)
        self.pushButton_reg.clicked.connect(self.button2_clicked)
        # # 使用样式表设置对话框的背景颜色  
        # Dialog.setStyleSheet("""  
        #     QDialog {  
        #         background-color: qlineargradient(x1: 0, y1: 0, x2: 0, y2: 1,  
        #                                           stop: 0 #333344, stop: 1 #333388);  
        #     }  
        # """)
        Dialog.setWindowTitle("人脸识别DEMO")
        self.label_imgshow.setText("")


    def button1_clicked(self):
        ui_faceadd = FaceAddDlg(Dialogfaceadd)
        Dialogfaceadd.exec()
    def button2_clicked(self):
        img_path,_ = QFileDialog.getOpenFileName(caption="选择图片文件")  
        imgtoreg = cv2.imread(img_path)
        if not imgtoreg is None:
            name_reg = "未识别"
            face_locations = face_recognition.face_locations(imgtoreg)
            if len(face_locations) == 1:
                show_img_encoding = face_recognition.face_encodings(imgtoreg)[0]
                #print("show_img_encoding:", len(show_img_encoding))
                face_distances = face_recognition.face_distance(g_list_faceencode, show_img_encoding)
                minvalue = min(face_distances)
                min_index = list(face_distances).index(minvalue)
                if minvalue < 0.4:
                    name_reg = g_list_facename[min_index]
                for i in range(len(face_locations)):
                    cv2.rectangle(imgtoreg,(face_locations[i][1],face_locations[i][0]),(face_locations[i][3],face_locations[i][2]),color=(0,255,0))
                imgtoreg = cv2ImgAddText(imgtoreg, name_reg, 10, 10)
                rgb_image = cv2.cvtColor(imgtoreg, cv2.COLOR_BGR2RGB) 
                h, w, ch = rgb_image.shape  
                bytes_per_line = ch * w  
                convert_to_Qt_format = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)  
                # # 将 QImage 转换为 QPixmap 并显示在 QLabel 上  
                # pixmap = QPixmap.fromImage(convert_to_Qt_format)  
                # 将 QImage 转换为 QPixmap 并保持宽高比  
                pixmap = QPixmap.fromImage(convert_to_Qt_format).scaled(  
                    self.label_imgshow.size(),  
                    Qt.KeepAspectRatio,  
                    Qt.SmoothTransformation  
                ) 
                # 将 QPixmap 设置为 QLabel 的内容  
                self.label_imgshow.setPixmap(pixmap) 

    

if __name__=="__main__":
    app = QtWidgets.QApplication(sys.argv)
    Dialog = QtWidgets.QDialog()
    Dialogfaceadd = QtWidgets.QDialog()
    ui_main = MainDlg(Dialog)
    Dialog.show()
    sys.exit(app.exec_())